import com.sun.corba.se.impl.encoding.CodeSetConversion;

/**
 * Creat with IntelliJ IDEA
 * Description:
 * User:WZW
 * Date:2022-10-13
 * Time:16:10
 */
public class RBTree {

    static class RBTreeNode {
        public int val;
        public RBTreeNode left;// 左子树的引用
        public RBTreeNode right;// 右子树的引用
        public RBTreeNode parent;// 父亲节点的引用
        public COLOR color = COLOR.RED;

        public RBTreeNode(int val) {
            this.val = val;
        }
    }

    public RBTreeNode root;

    // 假设有X个黑色节点，那么整颗红黑树当中的节点个数为N[X,2X]
    // 最短路径是logX这个是趋近logN
    // 查找的时间复杂度为O(logN)
    // 新增的节点不能是黑色的
    public boolean insert(int val) {
        RBTreeNode node = new RBTreeNode(val);
        if (root == null) {
            root = node;
            return true;
        }

        RBTreeNode parent = null;
        RBTreeNode cur = root;
        // 查找插入的位置
        while (cur != null) {
            if (cur.val < val) {
                parent = cur;
                cur = cur.right;
            } else if (cur.val == val) {
                return false;
            } else {
                parent = cur;
                cur = cur.left;
            }
        }
        if (parent.val < val) {
            parent.right = node;
        } else {
            parent.left = node;
        }
        node.parent = parent;
        cur = node;
        // 调节颜色
        // parent存在且为红色
        while (parent != null && parent.color == COLOR.RED) {
            RBTreeNode grandFather = parent.parent;
            if (parent == grandFather.left) {
                RBTreeNode uncle = grandFather.right;
                if (uncle != null && uncle.color == COLOR.RED) {
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    // 继续向上修改
                    cur = grandFather;
                    parent = cur.parent;
                } else {
                    //uncle不存在或uncle为黑
                    if (cur == parent.right) {
                        rotateLeft(parent);
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    } // 情况3 -> 情况2
                    //情况2
                    rotateRight(grandFather);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            } else {
                //parent == grandFather.right
                RBTreeNode uncle = grandFather.left;
                if (uncle != null && uncle.color == COLOR.RED) {
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    //向上继续修改
                    cur = grandFather;
                    parent = cur.parent;
                } else {
                    //uncle不存在或为黑色
                    //情况3
                    if (cur == parent.left) {
                        rotateRight(parent);
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    } //变为情况2
                    // 情况2
                    rotateLeft(parent);
                    grandFather.color = COLOR.RED;
                    parent.color = COLOR.BLACK;
                }
            }
        }
        return true;
    }

    public void rotateRight(RBTreeNode parent) {
        RBTreeNode subL = parent.left;
        RBTreeNode subLR = subL.right;

        parent.left = subLR;
        subL.right = parent;
        if (subLR != null) {
            subLR.parent = parent;
        }

        RBTreeNode pParent = parent.parent;
        parent.parent = subL;
        // 检查当前是不是根节点
        if (parent == root) {
            root = subL;
            root.parent = null;
        } else {
            // 不是根节点，判断这颗子树是左子树还是右子树
            if (pParent.left == parent) {
                pParent.left = subL;
            }else {
                pParent.right = subL;
            }
            subL.parent = parent.parent;
        }
        return;
    }

    public void rotateLeft(RBTreeNode parent) {
        RBTreeNode subR = parent.right;
        RBTreeNode subRL = subR.left;

        parent.right = subRL;
        if (subRL != null) {
            subRL.parent = parent;
        }
        subR.left = parent;

        if (parent == root) {
            root = subR;
            root.parent = null;
        } else {
            RBTreeNode pParent = parent.parent;
            if (pParent.left == parent) {
                pParent.left = subR;
            }else {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
        parent.parent = subR;
    }

    // 验证红黑树
    public boolean isBRTree() {
        if (root == null) return true;
        if (root.color == COLOR.RED) return false;
        // 存储当前红黑树中 最左边路径的黑色节点的个数
        int blackNum = 0;
        RBTreeNode cur = root;
        while (cur != null) {
            if (cur.color == COLOR.BLACK) {
                blackNum++;
            }
            cur = cur.left;
        }
        //检查是否存在两个连续的红色节点 && 每条路径上黑色节点的个数是一直的
        return checkRedColor(root) && checkBlackNum(root, 0, blackNum);
    }

    /**
     *
     * @param root
     * @param pathBlackNum 每次递归，计算黑色节点的个数
     * @param blackNum 事先计算好的某条路径上的黑色节点的个数
     * @return
     */
    private boolean checkBlackNum(RBTreeNode root, int pathBlackNum, int blackNum) {
        if (root == null) return true;
        if (root.color == COLOR.BLACK) {
            pathBlackNum++;
        }
        if (root.left == null && root.right == null) {
            if (pathBlackNum != blackNum)
                return false;
        }
        return checkBlackNum(root.left, pathBlackNum, blackNum) && checkBlackNum(root.right, pathBlackNum, blackNum);
    }

    // 检查是否存在两个连续的红色节点，不存在返回true
    private boolean checkRedColor(RBTreeNode root) {
        if (root == null) return true;
        // 如果出现红色节点，判断它的父亲是不是红色节点
        if (root.color == COLOR.RED) {
            RBTreeNode parent = root.parent;
            if (parent.color == COLOR.RED) //连续出现两个红色节点
                return false;
        }
        // 递归遍历
        return checkRedColor(root.left) && checkRedColor(root.right);
    }
}
